import numpy as np
import random
import cv2
import os
import os.path as osp
import shutil
import json


def set_seed(manual_seed):
    random.seed(manual_seed)  
    np.random.seed(manual_seed)  


seed = 7
set_seed(seed)
copy_times = 5

BASEDIR = osp.dirname(osp.abspath(__file__))


IMGDIR = osp.join(BASEDIR, 'train')
LABELDIR = osp.join(BASEDIR, 'train')
OUTDIR = osp.join(BASEDIR, f'cpAug_AGH_train_seed{seed}')

# IMGDIR = osp.join(BASEDIR, 'val')
# LABELDIR = osp.join(BASEDIR, 'val')
# OUTDIR = osp.join(BASEDIR, f'cpAug_AGH_val_seed{seed}')

if osp.exists(OUTDIR):
    shutil.rmtree(OUTDIR)
os.makedirs(OUTDIR)


label_to_num_dict = {
    'A': 0,
    'B': 1,
    'C': 2,
    'D': 3,
    'E': 4,
    'F': 5,
    'G': 6,
    'H': 7,
    'I': 8,
    'J': 9,
    'K': 10,
}

num_to_name_dict = {}
for k in label_to_num_dict.keys():
    num_to_name_dict[label_to_num_dict[k]] = k


class SmallObjectAugmentation(object):
    def __init__(self, thresh=64*64, prob=1.0, copy_times=copy_times, epochs=30, all_objects=False, one_object=False):
        """
        sample = {'img':img, 'annot':annots}
        img = [height, width, 3]
        annot = [xmin, ymin, xmax, ymax, label]
        thresh：the detection threshold of the small object. If annot_h * annot_w < thresh, the object is small
        prob: the prob to do small object augmentation
        epochs: the epochs to do
        """
        self.thresh = thresh
        self.prob = prob
        self.copy_times = copy_times
        self.epochs = epochs
        self.all_objects = all_objects
        self.one_object = one_object
        if self.all_objects or self.one_object:
            self.copy_times = 1

    def issmallobject(self, h, w):
        if h * w <= self.thresh:
            return True
        else:
            return False

    def compute_overlap(self, annot_a, annot_b):
        if annot_a is None: return False
        left_max = max(annot_a[0], annot_b[0])
        top_max = max(annot_a[1], annot_b[1])
        right_min = min(annot_a[2], annot_b[2])
        bottom_min = min(annot_a[3], annot_b[3])
        inter = max(0, (right_min-left_max)) * max(0, (bottom_min-top_max))
        if inter != 0:
            return True
        else:
            return False

    def donot_overlap(self, new_annot, annots):
        for annot in annots:
            if self.compute_overlap(new_annot, annot): return False
        return True

    def create_copy_annot(self, h, w, annot, annots, ro_w, ro_h, angle):
        annot = annot.astype(np.int)
        annot_h, annot_w = annot[3] - annot[1], annot[2] - annot[0]
        for epoch in range(self.epochs):
            random_x, random_y = np.random.randint(int(annot_w / 2), int(w - annot_w / 2)), \
                                 np.random.randint(int(annot_h / 2), int(h - annot_h / 2))
            xmin, ymin = random_x - annot_w / 2, random_y - annot_h / 2
            xmax, ymax = xmin + annot_w, ymin + annot_h
            if xmin < 0 or xmax > w or ymin < 0 or ymax > h:
                continue
            new_annot = np.array([xmin, ymin, xmax, ymax, annot[4], ro_w, ro_h, angle]).astype(np.int)

            if self.donot_overlap(new_annot, annots) is False:
                continue

            return new_annot
        return None

    def add_patch_in_img(self, annot, copy_annot, image):
        copy_annot = copy_annot.astype(np.int)
        image[annot[1]:annot[3], annot[0]:annot[2], :] = image[copy_annot[1]:copy_annot[3], copy_annot[0]:copy_annot[2], :]
        return image

    def __call__(self, sample):
        
        # sample: dict
        # sample['img']: np.ndarray
        # sample['annot']: -> top left, bottom right
        
        if self.all_objects and self.one_object: return sample
        if np.random.rand() > self.prob: return sample

        img, annots = sample['img'], sample['annot']
        h, w = img.shape[0], img.shape[1]
        
        small_object_list = list()
        for idx in range(annots.shape[0]):
            # annot = annots[idx]
            # annot_h, annot_w = annot[3] - annot[1], annot[2] - annot[0]
            # if self.issmallobject(annot_h, annot_w):
            if True:  # 都当成小目标
                small_object_list.append(idx)
        # print('small_object_list:', small_object_list)

        l = len(small_object_list)
        # No Small Object
        if l == 0: return sample

        # Refine the copy_object by the given policy
        # Policy 2:
        # copy_object_num = np.random.randint(0, l)
        copy_object_num = l  # 全都复制粘贴
        
        # print(l)
        # print(copy_object_num)
        
        # Policy 3:
        if self.all_objects:
            copy_object_num = l
            print('Policy 3')
            
        # Policy 1:
        if self.one_object:
            copy_object_num = 1
            print('Policy 1')

        random_list = random.sample(range(l), copy_object_num)
        # print('random_list:', random_list)
        
        annot_idx_of_small_object = [small_object_list[idx] for idx in random_list]
        # print('annot_idx_of_small_object:', annot_idx_of_small_object)
        
        select_annots = annots[annot_idx_of_small_object, :]
        # print('select_annots:', select_annots)
        
        # 按类别筛选
        select_annots = [item for item in select_annots if int(item[4]) in (label_to_num_dict['A'], label_to_num_dict['G'], label_to_num_dict['H'])]
        
        annots = annots.tolist()
        # for idx in range(copy_object_num):
        for idx in range(len(select_annots)):
            annot = select_annots[idx]
            # annot_h, annot_w = annot[3] - annot[1], annot[2] - annot[0]

            # if self.issmallobject(annot_h, annot_w) is False: continue
            
            ro_w, ro_h, angle = annot[-3:]
            
            for i in range(self.copy_times):
                new_annot = self.create_copy_annot(h, w, annot, annots, ro_w, ro_h, angle)
                if new_annot is not None:
                    img = self.add_patch_in_img(new_annot, annot, img)
                    annots.append(new_annot)

        return {'img': img, 'annot': np.array(annots)}


# start Augmentation processing...
imgnames = [name for name in os.listdir(IMGDIR) if name.endswith('.png')]
for imgname in imgnames:
    print(imgname, end='  ')
    jsonname = imgname.split('.')[-2]+'.json'
    try:
        with open(osp.join(LABELDIR, jsonname), 'r') as fp:  # 找到对应的json标注文件
            js_data = json.load(fp)
    except:
        print('json read bug, continue')
        continue
    boxes = js_data['shapes']
    
    has_special = False
    for box in boxes:
        cls_name = box['label']
        # print(cls_name)
        if cls_name in ["A", "G", "H"]:
            has_special = True
            break
        
    if not has_special:
        print('continue')
        continue
    
    path_img = osp.join(IMGDIR, imgname)
    img = cv2.imread(path_img)  # 读图
    
    annots = []
    for box in boxes:
        cls_name = box['label']
        
        xy4 = box['points']
        xy4 = np.array(xy4).astype(np.int0)  # 损失精度!
        ro_rect = cv2.minAreaRect(xy4)
        (cx,cy),(ro_w,ro_h),angle = ro_rect
        xmin = min(xy4[:, 0])
        xmax = max(xy4[:, 0])
        ymin = min(xy4[:, 1])
        ymax = max(xy4[:, 1])
        label = label_to_num_dict[cls_name]
        annot = [xmin, ymin, xmax, ymax, label, ro_w, ro_h, angle]
        annots.append(annot)

    annots = np.array(annots, dtype=np.float32)
    # print(annots)
    
    sample_pair = {}
    sample_pair['img'] = img
    sample_pair['annot'] = annots
    
    # 开始执行cut-paste数据增强
    print('do augmentation')
    print(annots.shape)  # 统计增强前目标个数
    aug = SmallObjectAugmentation()
    try:
        ret = aug(sample_pair)
    except:
        print('do aug bug, continue')
        continue
    # print(ret)
    new_img = ret['img']
    new_ann = ret['annot']
    print(new_ann.shape)  # 统计增强后目标个数
    
    # save new image file
    new_imgname = f"{imgname.split('.')[-2]}_cp{seed}.png"
    path_img_out = osp.join(OUTDIR, new_imgname)
    cv2.imwrite(path_img_out, new_img)
    
    # save new annotation file
    boxes = []
    for i in range(new_ann.shape[0]):
        xmin, ymin, xmax, ymax, cls_num, ro_w, ro_h, angle = new_ann[i]
        
        cls_name = num_to_name_dict[cls_num]
        # print(cls_name)
        
        cx = (xmin + xmax) / 2
        cy = (ymin + ymax) / 2
        ro_rect = ((cx, cy), (ro_w, ro_h), angle)
        xy4 = cv2.boxPoints(ro_rect).tolist()  # shape:(4, 2) dtype:float32
        # print(xy4)
        
        box = {}
        box['label'] = cls_name
        box['points'] = xy4
        box["group_id"] = None,
        box["shape_type"] = "polygon",
        box["flags"] = {}

        boxes.append(box)
    
    labelData = {}
    labelData['shapes'] = boxes
    labelData["version"] = "4.5.6",
    labelData["flags"] = {},
    labelData["imagePath"] = new_imgname,
    labelData["imageHeight"] = new_img.shape[0],
    labelData["imageWidth"] = new_img.shape[1],
    labelData["imageData"] = None

    new_jsonname = new_imgname.split('.')[-2]+'.json'
    path_json_out = osp.join(OUTDIR, new_jsonname)
    with open(path_json_out, 'w') as fp:    
        json.dump(labelData, fp, indent=2)
    print(imgname, " done.")
